import os
import sys
from PIL import Image
import shutil

def process_images(source_dir, color_dir, bw_dir, prefix_to_remove=""):
    """
    处理图片：去掉前缀、转换为PNG、生成黑白版本
    
    Args:
        source_dir (str): 源图片目录路径
        color_dir (str): 彩色PNG图片保存目录
        bw_dir (str): 黑白图片保存目录
        prefix_to_remove (str): 要去除的文件名前缀
    """
    # 检查源目录是否存在
    if not os.path.exists(source_dir):
        print(f"错误: 源目录 '{source_dir}' 不存在")
        return False
    
    # 创建目标目录
    os.makedirs(color_dir, exist_ok=True)
    os.makedirs(bw_dir, exist_ok=True)
    
    # 支持的图片格式
    supported_formats = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp', '.gif'}
    
    # 统计信息
    processed_count = 0
    error_count = 0
    
    print("开始处理图片...")
    
    # 遍历源目录中的所有文件
    for filename in os.listdir(source_dir):
        # 检查文件是否为图片
        file_ext = os.path.splitext(filename)[1].lower()
        if file_ext not in supported_formats:
            continue
        
        # 去除指定前缀
        new_filename = filename
        if prefix_to_remove and filename.startswith(prefix_to_remove):
            new_filename = filename[len(prefix_to_remove):]
        
        # 确保新文件名以.png结尾
        new_filename = os.path.splitext(new_filename)[0] + ".png"
        
        color_path = os.path.join(color_dir, new_filename)
        bw_path = os.path.join(bw_dir, new_filename)
        
        source_path = os.path.join(source_dir, filename)
        
        try:
            # 打开原始图片
            with Image.open(source_path) as img:
                # 转换为RGB模式（如果图片是RGBA或P模式）
                if img.mode in ('RGBA', 'P'):
                    img = img.convert('RGB')
                
                # 保存彩色PNG版本
                img.save(color_path, 'PNG')
                
                # 转换为黑白并保存
                bw_img = img.convert('L')
                bw_img.save(bw_path, 'PNG')
                
                print(f"处理成功: {filename} -> {new_filename}")
                processed_count += 1
                
        except Exception as e:
            print(f"处理图片失败 {filename}: {str(e)}")
            error_count += 1
    
    # 输出统计信息
    print(f"\n处理完成!")
    print(f"成功处理: {processed_count} 张图片")
    print(f"失败: {error_count} 张图片")
    print(f"彩色PNG图片保存在: {color_dir}")
    print(f"黑白PNG图片保存在: {bw_dir}")
    
    return True

def main():
    """主函数"""
    print("=== 图片批量处理工具 ===")
    print("功能: 1. 去掉文件名前缀 2. 转换为PNG格式 3. 生成黑白版本")
    
    # 获取用户输入
    source_dir = "./baidu_images"
   
    
    color_dir = "./baidu_images_color"
    
    bw_dir = "./baidu_images_bw"
    
    prefix_to_remove = "全彩漫画_"
    
    # 确认操作
    print(f"\n即将执行以下操作:")
    print(f"1. 处理目录: {source_dir}")
    print(f"2. 彩色PNG保存到: {color_dir}")
    print(f"3. 黑白PNG保存到: {bw_dir}")
    if prefix_to_remove:
        print(f"4. 去除文件名前缀: '{prefix_to_remove}'")
    
    # 执行处理
    success = process_images(source_dir, color_dir, bw_dir, prefix_to_remove)
    
    if success:
        print("所有操作已完成!")
    else:
        print("操作过程中出现问题!")

if __name__ == "__main__":
    main()